package com.newbee.aop;

import com.newbee.model.MqCallbackBaseMessage;
import com.newbee.utils.TraceContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;

/**
 * MQ回调专用AOP：自动从任意位置的参数中解析traceId并设置到线程上下文
 */
@Aspect
@Component
@Slf4j
public class MqCallbackTraceAspect {

    // 切入点：拦截所有MQ回调相关的Controller方法（可根据实际包路径调整）
    @Pointcut("execution(* com.newbee.controller.MockMqCallbackController.receiveCallback(..))")
    public void mqCallbackPointcut() {
    }

    /**
     * 前置通知：循环查找参数中的MQ消息对象，解析并设置traceId
     */
    @Before("mqCallbackPointcut()")
    public void beforeMqCallback(JoinPoint joinPoint) {
        // 1. 获取方法所有参数
        Object[] args = joinPoint.getArgs();
        if (args == null || args.length == 0) {
            log.warn("MQ回调方法无参数，无法解析traceId");
            return;
        }

        // 2. 循环查找第一个MqCallbackBaseMessage类型的参数
        MqCallbackBaseMessage message = findFirstMqMessage(args);
        if (message == null) {
            log.warn("MQ回调方法参数中未找到MqCallbackBaseMessage类型的对象");
            return;
        }

        // 3. 解析并设置traceId到上下文
        String traceId = message.getTraceId();
        String asyncRootNodeId = message.getAsyncRootNodeId();
        if (StringUtils.isAnyBlank(traceId, asyncRootNodeId)) {
            log.warn("MQ回调消息中traceId或asyncRootNodeId为空");
            return;
        }

        TraceContextHolder.initForAsync(traceId, asyncRootNodeId);
        log.info("MQ回调AOP自动设置traceId：{} {}（从参数中解析）", traceId, asyncRootNodeId);
    }

    /**
     * 循环查找参数列表中第一个MqCallbackBaseMessage类型（或子类）的对象
     */
    private MqCallbackBaseMessage findFirstMqMessage(Object[] args) {
        for (Object arg : args) {
            // 支持子类（如继承MqCallbackBaseMessage的业务消息类）
            if (arg instanceof MqCallbackBaseMessage) {
                return (MqCallbackBaseMessage) arg;
            }
        }
        return null;
    }
}